package com.eckey.lab.controller;

import lombok.extern.slf4j.Slf4j;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @Author: ChengLiang
 * @CreateTime: 2023-10-08  14:29
 * @Description: TODO
 * @Version: 1.0
 */
@Slf4j
@RestController
@CrossOrigin
@RequestMapping("/sse")
public class SSEController {

    /**
     * 使用Map对象来存放userId对应的
     */
    private static Map<String, SseEmitter> sseEmitterMap = new ConcurrentHashMap<>();

    @GetMapping(path = "/subscribe/{userId}", produces = {MediaType.TEXT_EVENT_STREAM_VALUE})
    public SseEmitter subscribe(@PathVariable("userId") String userId) throws IOException {
        // 设置超时时间，0表示不过期。默认30秒，超过时间未完成会抛出异常：AsyncRequestTimeoutException
        SseEmitter sseEmitter = new SseEmitter(30_000L);
        // 设置前端的重试时间为15s，如果不加这个发送一下，前端就不会显示连接成功
        sseEmitter.send("连接成功");
        // 注册回调
        sseEmitter.onCompletion(() -> {
            log.info("连接结束：{}", userId);
        });

        sseEmitter.onError((Throwable throwable) -> {
            log.error("连接异常:{}", throwable.getMessage());
        });
        sseEmitter.onTimeout(() -> {
            log.warn("连接超时：{}", userId);
        });
        //以userId为key，如果一个用户多个设备连接，会不准确
        sseEmitterMap.put(userId, sseEmitter);
        log.info("创建新的sse连接，当前用户：{}", userId);
        return sseEmitter;
    }

    @GetMapping(path = "/sendMessage")
    public void sendMessage(String userId, String message) {
        if (sseEmitterMap.containsKey(userId)) {
            try {
                log.info("向用户：{},发送消息：{}", userId, message);
                sseEmitterMap.get(userId).send(message, MediaType.APPLICATION_JSON);
            } catch (IOException e) {
                log.error("用户[{}]推送异常:{}", userId, e.getMessage());
                removeUser(userId);
            }
        }
    }

    private void removeUser(String userId) {
        sseEmitterMap.remove(userId);
        log.info("移除用户成功：{}", userId);
    }

    @GetMapping("/close/{userId}")
    public void close(@PathVariable("userId") String userId) {
        removeUser(userId);
        log.info("关闭连接成功：{}", userId);
    }
}
